import asyncio
import os
import datetime

from sqlalchemy import exc, select
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from stqdm import stqdm

from models.models import FileForecast, Forecast
from tasks.task_manager import task_queue


def first_forecast_date(date, num_week):
    wd = date.isocalendar()[2]
    # dist = 22 - wd if wd > 1 else 21
    dist = num_week * 7 + 1 - wd if wd > 1 else num_week * 7
    return date + datetime.timedelta(days=dist)


def gen_forecast_dates(from_date, num_weeks):
    return [from_date + datetime.timedelta(days=i * 7) for i in range(num_weeks)]


def season_fit(x, y):
    vec_len = x.unique().max() + 1
    res = dict()
    for i in range(vec_len):
        res[i] = round(y.loc[x[x == i].index].mean())
    return res


def target_model(work, template, modality, modality_type):
    df = template.copy()
    df["modality"] = modality
    df["type"] = modality_type
    df["kind"] = "Прогноз"

    xcol0 = ["wc"]
    ycol = "number"
    model = LinearRegression().fit(work[xcol0], work[ycol])
    work["d0"] = work.number - model.predict(work[xcol0])
    v = season_fit(work["wn"], work["d0"])
    work["d1"] = work["wn"].map(v)
    df["d1"] = df["wn"].map(v)
    xcol1 = ["wc", "d1"]
    model.fit(work[xcol1], work[ycol])
    df["number"] = np.round(model.predict(df[xcol1])).astype(np.int32)
    return df.drop(columns=["wn", "d1"])


async def task_create_forecast(statistics_filename: str, forecast_name: str):
    """Задача формирования прогноза."""
    SQLALCHEMY_DATABASE_URL = "postgresql+asyncpg://user:pass@postgres_db:5432/dbname"
    engine = create_async_engine(SQLALCHEMY_DATABASE_URL)

    async_session = sessionmaker(
        engine,
        class_=AsyncSession,
        expire_on_commit=False,
    )

    task_queue.put(
        item={"task_type": "forecast", "task_name": forecast_name, "status": "DOING"}
    )
    async with async_session() as session:
        query = select(FileForecast).where(FileForecast.filename == forecast_name)
        query_result = await session.execute(query)
        forecast_file = query_result.scalars().first()
        if forecast_file:
            task_queue.put(
                item={
                    "task_type": "forecast",
                    "task_name": forecast_name,
                    "status": "ERROR",
                    "error": "Файл уже существует в базе данных",
                }
            )

        forecast_file = FileForecast(filename=forecast_name)
        session.add(forecast_file)
        try:
            await session.commit()
            await session.refresh(forecast_file)
        except (exc.IntegrityError, RuntimeError):
            await session.rollback()
            task_queue.put(
                item={
                    "task_type": "forecast",
                    "task_name": forecast_name,
                    "status": "ERROR",
                    "error": "Файл уже существует в базе данных",
                }
            )
            return

    num_week_step = 1
    num_weeks = 14
    data = pd.read_csv(
        os.path.join("download", statistics_filename), parse_dates=["date"]
    )
    from_date = first_forecast_date(data.date.max(), num_week_step)
    from_week_seq = data.wc.max() + num_week_step

    task_queue.put(
        {
            "task_type": "forecast",
            "task_name": forecast_name,
            "progress": 1,
            "msg": "Начало загрузки данных...",
        }
    )

    if "wn" not in data.columns:
        data["wn"] = data.date.apply(lambda x: x.isocalendar()[1] - 1)

    norm = pd.read_csv(os.path.join("download", "Labour_norm.csv"))
    norm.columns = ["modality", "-", "КУ1", "КУ2+"]

    ft = pd.DataFrame(
        np.array(gen_forecast_dates(from_date, num_weeks)), columns=["date"]
    )
    ft["year"] = ft.date.apply(lambda x: x.year)
    ft["week"] = ft.date.apply(lambda x: x.isocalendar()[1])
    ft["wn"] = ft.week - 1
    ft["wc"] = np.arange(len(ft)) + from_week_seq

    # target = 'number'
    # features = ['wn']

    for_list = []

    task_queue.put(
        {
            "task_type": "forecast",
            "task_name": forecast_name,
            "progress": 2,
            "msg": "Данные загружены.",
        }
    )

    task_queue.put(
        {
            "task_type": "forecast",
            "task_name": forecast_name,
            "progress": 4,
            "msg": "Начало обучения модели...",
        }
    )

    progress = 10

    mod_list = list(data.modality.unique())
    for mi in stqdm(mod_list):
        mod_sel = data[data.modality == mi]
        norm_sel = norm[norm.modality == mi]
        sel_list = list(mod_sel.type.unique())
        for si in sel_list:
            df = target_model(mod_sel[mod_sel.type == si].copy(), ft, mi, si)
            whnorm = norm_sel[si].iloc[0]
            df["workhours"] = df.number.apply(lambda x: round(x * whnorm / 60))
            for_list.append(df)
        task_queue.put(
            {
                "task_type": "forecast",
                "task_name": forecast_name,
                "progress": progress,
                "msg": f"Обработано {mi} из {len(mod_list)} модальностей.",
            }
        )

        progress += 10
    task_queue.put(
        {
            "task_type": "forecast",
            "task_name": forecast_name,
            "progress": 70,
            "msg": "Завершение обучения модели.",
        }
    )
    forecast = pd.concat(for_list, axis=0).reset_index(drop=True)

    task_queue.put(
        {
            "task_type": "forecast",
            "task_name": forecast_name,
            "progress": 80,
            "msg": "Начало сохранения прогноза...",
        }
    )

    # forecast.to_csv(forecast_name, index=False)
    records = forecast.to_dict("records")
    # return forecast.shape[0]
    async with async_session() as session:
        for record in records:
            forecast_record = Forecast(
                id_filename=forecast_file.id,
                date=record["date"].date(),
                week=record["week"],
                wc=record["wc"],
                modality=record["modality"],
                type=record["type"],
                kind=record["kind"],
                number=record["number"],
                workhours=record["workhours"],
            )
            session.add(forecast_record)
        await session.commit()

    task_queue.put(
        {
            "task_type": "forecast",
            "task_name": forecast_name,
            "progress": 100,
            "msg": "Прогноз успешно сохранен.",
            "status": "DONE",
        }
    )


def start_forecast_task(
    statistics_filename: str,
    forecast_name: str,
):
    """Запускает асинхронную функцию в отдельном потоке."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(
        task_create_forecast(
            statistics_filename=statistics_filename,
            forecast_name=forecast_name,
        )
    )
